import sqlalchemy
from sqlalchemy import create_engine, Column, Date, VARCHAR, FLOAT, Integer, TIMESTAMP, func, Index, DECIMAL
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import Table

"""
    创建ORM映射对象
"""

Base = declarative_base()

engine = create_engine(
    'postgresql+psycopg2://postgres:941010@localhost:5432/quant_analysis',
    echo=False, pool_size=20,
    connect_args={'client_encoding': 'utf8', 'options': '-csearch_path=quant_research'},
    isolation_level='AUTOCOMMIT',
    pool_pre_ping=True, pool_recycle=3600
)

engine2 = sqlalchemy.engine.create_engine(
    "postgresql://root:Ct941010$@114.215.186.147/stockselectors",
    echo=False, pool_size=20, connect_args={'client_encoding': 'utf8'},
    isolation_level='AUTOCOMMIT',
    pool_pre_ping=True, pool_recycle=3600
)

Session = sessionmaker(bind=engine)


class StockInfo(Base):
    __tablename__ = 'basic_info_stock'
    __table_args__ = {'schema': 'quant_research'}

    ticker = Column(VARCHAR, primary_key=True)
    short_name = Column(VARCHAR)
    area = Column(VARCHAR)
    industry = Column(VARCHAR)
    market = Column(VARCHAR)
    list_date = Column(Date)
    act_name = Column(VARCHAR)
    act_ent_type = Column(VARCHAR)
    status = Column(VARCHAR)

    def __repr__(self):
        return "<basic_info_stock(name=basic_info_stock, comment=A股基本信息)>"


class TsMarketDaily(Base):
    __tablename__ = 'market_daily_ts'
    __table_args__ = {'schema': 'quant_research'}
    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    ticker = Column(VARCHAR)
    open = Column(FLOAT)
    close = Column(FLOAT)
    pre_close = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    change = Column(FLOAT)
    pct_chg = Column(FLOAT)
    amount = Column(FLOAT)
    vol = Column(FLOAT)

    def __repr__(self):
        return "<MarketDaily(name=marketData_daily_ts, comment=tushare行情接口记录全A日线行情数据)>"


class Market1Min(Base):
    __tablename__ = 'market_1min'
    __table_args__ = {'schema': 'quant_research'}
    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_time = Column(TIMESTAMP)
    ticker = Column(VARCHAR)
    open = Column(FLOAT)
    close = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    amount = Column(FLOAT)
    volumn = Column(FLOAT)

    def __repr__(self):
        return "<Market1min(name=market_1min, comment=记录全A分钟线行情数据)>"


class TsIndicator(Base):
    __tablename__ = 'indicator_daily'
    __table_args__ = {'schema': 'quant_research'}
    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    code = Column(VARCHAR)
    close = Column(FLOAT)
    turnover_rate = Column(FLOAT)
    turnover_rate_f = Column(FLOAT)
    volume_ratio = Column(FLOAT)
    pe = Column(FLOAT)
    pe_ttm = Column(FLOAT)
    pb = Column(FLOAT)
    ps = Column(FLOAT)
    ps_ttm = Column(FLOAT)
    dv_ratio = Column(FLOAT)
    dv_ttm = Column(FLOAT)
    total_share = Column(FLOAT)
    float_share = Column(FLOAT)
    free_share = Column(FLOAT)
    total_mv = Column(FLOAT)
    circ_mv = Column(FLOAT)

    def __repr__(self):
        return "<TsIndicator(name=indicator_daily, comment=记录全A个股指标日行情数据)>"


# 个股基本信息数据格式字典
dtype_basic = {
    'ticker': sqlalchemy.types.VARCHAR, 'short_name': sqlalchemy.types.VARCHAR, 'area': sqlalchemy.types.VARCHAR,
    'industry': sqlalchemy.types.VARCHAR, 'market': sqlalchemy.types.VARCHAR, 'list_date': sqlalchemy.types.DATE,
    'act_name': sqlalchemy.types.VARCHAR, 'act_ent_type': sqlalchemy.types.VARCHAR, 'status': sqlalchemy.types.VARCHAR,
}

# ts日线行情
dtype_ts_bardaily = {
    'ticker': sqlalchemy.types.VARCHAR, 'trade_date': sqlalchemy.types.DATE, 'open': sqlalchemy.types.FLOAT,
    'high': sqlalchemy.types.FLOAT, 'low': sqlalchemy.types.FLOAT, 'close': sqlalchemy.types.FLOAT,
    'pre_close': sqlalchemy.types.FLOAT, 'change': sqlalchemy.types.FLOAT, 'pct_chg': sqlalchemy.types.FLOAT,
    'vol': sqlalchemy.types.FLOAT, 'amount': sqlalchemy.types.FLOAT,
}

# ts指数日行情
dtype_ts_indexdaily = {
    'index_code': sqlalchemy.types.VARCHAR, 'trade_date': sqlalchemy.types.DATE, 'close': sqlalchemy.types.FLOAT,
    'open': sqlalchemy.types.FLOAT, 'high': sqlalchemy.types.FLOAT, 'low': sqlalchemy.types.FLOAT,
    'pre_close': sqlalchemy.types.FLOAT, 'change': sqlalchemy.types.FLOAT, 'pct_chg': sqlalchemy.types.FLOAT,
    'vol': sqlalchemy.types.FLOAT, 'amount': sqlalchemy.types.FLOAT,
}

dtype_indicators = {
    'trade_date': sqlalchemy.types.DATE, 'code': sqlalchemy.types.VARCHAR, 'close': sqlalchemy.types.FLOAT,
    'turnover_rate': sqlalchemy.types.FLOAT, 'turnover_rate_f': sqlalchemy.types.FLOAT,
    'volume_ratio': sqlalchemy.types.FLOAT,
    'pe': sqlalchemy.types.FLOAT, 'pe_ttm': sqlalchemy.types.FLOAT, 'pb': sqlalchemy.types.FLOAT,
    'ps': sqlalchemy.types.FLOAT,
    'ps_ttm': sqlalchemy.types.FLOAT, 'dv_ratio': sqlalchemy.types.FLOAT, 'dv_ttm': sqlalchemy.types.FLOAT,
    'total_share': sqlalchemy.types.FLOAT,
    'float_share': sqlalchemy.types.FLOAT, 'free_share': sqlalchemy.types.FLOAT, 'total_mv': sqlalchemy.types.FLOAT,
    'circ_mv': sqlalchemy.types.FLOAT,
}

"""
    构建因子数据类
"""


# 区间波动率因子值
class FactorPeriodVolatility(Base):
    __tablename__ = 'factor_periodVolatility'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trading = Column(Date)
    signal = Column(FLOAT)

    createtime = Column(TIMESTAMP, default=func.now())  # 自动设置创建时间
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())  # 自动更新数据时间

    __table_args__ = (
        Index('idx_f_periodVolatility', 'code', 'trading', unique=True),
    )


# 区间累计收益率因子值
class FactorReturnWgt(Base):
    __tablename__ = 'factor_returnWgt'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trading = Column(Date)
    signal = Column(FLOAT)
    pct_chg = Column(FLOAT)
    wgt_chg = Column(FLOAT)

    createtime = Column(TIMESTAMP, default=func.now())  # 自动设置创建时间
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())  # 自动更新数据时间

    __table_args__ = (
        Index('idx_f_returnWgt', 'code', 'trading', unique=True),
    )


# PE估值因子
class FactorPE(Base):
    __tablename__ = 'factor_ep_compound'

    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    code = Column(VARCHAR)
    ep = Column(FLOAT)
    end_date = Column(Date)
    is_open = Column(FLOAT)
    pretrade_date = Column(Date)
    n_income_attr_p = Column(FLOAT)
    income_attr_ttm = Column(FLOAT)
    open = Column(FLOAT)
    close = Column(FLOAT)
    ep_quantile = Column(FLOAT)
    ep_pivot_1y = Column(FLOAT)
    ep_pivot_3y = Column(FLOAT)

    __table_args__ = (
        Index('idx_ep', 'code', 'trade_date', unique=False),
    )


# PB估值因子
class FactorPB(Base):
    __tablename__ = 'factor_bp_compound'

    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    code = Column(VARCHAR)
    bp = Column(FLOAT)
    bp_quantile = Column(FLOAT)
    bp_pivot_1y = Column(FLOAT)
    bp_pivot_3y = Column(FLOAT)

    __table_args__ = (
        Index('idx_bp', 'code', 'trade_date', unique=False),
    )


# 成交量加权收益率偏度因子值
class FactorVolskew(Base):
    __tablename__ = 'factor_volSkew'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trading = Column(Date)
    signal_volReturn = Column(FLOAT)

    createtime = Column(TIMESTAMP, default=func.now())  # 自动设置创建时间
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())  # 自动更新数据时间

    __table_args__ = (
        Index('idx_f_volSkew', 'code', 'trading', unique=True),
    )


# RVI波动率因子
class FactorRVI(Base):
    __tablename__ = 'factor_rvi'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trading = Column(Date)
    open = Column(FLOAT)
    high = Column(FLOAT)
    close = Column(FLOAT)
    low = Column(FLOAT)
    riv = Column(FLOAT)
    riv_slope = Column(FLOAT)
    riv_accel = Column(FLOAT)
    riv_up = Column(FLOAT)
    riv_up_slope = Column(FLOAT)
    riv_down = Column(FLOAT)
    riv_down_slope = Column(FLOAT)

    __table_args__ = (
        Index('idx_f_rvi', 'code', 'trading', unique=True),
    )


# 指数成分
class IndexConstituent(Base):
    __tablename__ = 'index_constituent'
    id = Column(Integer, primary_key=True, autoincrement=True)
    index_code = Column(VARCHAR)
    con_code = Column(VARCHAR)
    trade_date = Column(Date)
    weight = Column(FLOAT)


# 业绩超预期因子（预测方法：随机游走）
class FactorEarningSuprise(Base):
    __tablename__ = 'factor_earningSuprise'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    f_ann_date = Column(Date)
    end_date = Column(Date)
    sue_profit = Column(FLOAT)
    sue_revenue = Column(FLOAT)

    __table_args__ = (
        Index('idx_f_earningSuprise', 'code', 'end_date', unique=True),
    )


# fama 三因子
class FactorFama(Base):
    __tablename__ = 'factor_fama'

    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    smb = Column(FLOAT)
    hml = Column(FLOAT)
    smb_small = Column(FLOAT)
    smb_large = Column(FLOAT)
    hml_small = Column(FLOAT)
    hml_large = Column(FLOAT)


# 分钟行情数据表
class MinuteBars(Base):
    __tablename__ = 'minute_bars'

    trading = Column(TIMESTAMP, primary_key=True)
    code = Column(VARCHAR, primary_key=True)
    open = Column(DECIMAL)
    close = Column(DECIMAL)
    high = Column(DECIMAL)
    low = Column(DECIMAL)
    volumn = Column(DECIMAL)
    amount = Column(DECIMAL)


# 指数日行情数据类
class IndexDaily(Base):
    __tablename__ = 'index_daily'

    id = Column(Integer, primary_key=True, autoincrement=True)
    ts_code = Column(VARCHAR)
    trade_date = Column(Date)
    close = Column(FLOAT)
    open = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    pre_close = Column(FLOAT)
    change = Column(FLOAT)
    pct_chg = Column(FLOAT)
    vol = Column(FLOAT)
    amount = Column(FLOAT)
    index_name = Column(VARCHAR)

    __table_args__ = (
        Index('idx_indexdaily', 'trade_date', 'ts_code', 'index_name', unique=True),
    )


# 个股日行情数据类(无复权)
class StockDaily(Base):
    __tablename__ = 'market_daily_none'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trade_date = Column(Date)
    open = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    close = Column(FLOAT)
    pre_close = Column(FLOAT)
    change = Column(FLOAT)
    pct_chg = Column(FLOAT)
    vol = Column(FLOAT)
    amount = Column(FLOAT)

    __table_args__ = (
        Index('idx_stockdaily_none', 'trade_date', 'code', unique=True),
    )


# 个股日行情数据类(后复权)
class StockDailyHFQ(Base):
    __tablename__ = 'market_daily_ts'

    id = Column(Integer, primary_key=True, autoincrement=True)
    ticker = Column(VARCHAR)
    trade_date = Column(Date)
    open = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    close = Column(FLOAT)
    pre_close = Column(FLOAT)
    change = Column(FLOAT)
    pct_chg = Column(FLOAT)
    vol = Column(FLOAT)
    amount = Column(FLOAT)
    createtime = Column(TIMESTAMP, default=func.now())  # 自动设置创建时间
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())  # 自动更新数据时间

    __table_args__ = (
        Index('idx_stockdaily_hfq', 'ticker', 'trade_date', unique=True),
    )


# 个股日行情数据类（前复权）
class StockDailyQFQ(Base):
    __tablename__ = 'market_daily_qfq'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    trade_date = Column(Date)
    open = Column(FLOAT)
    high = Column(FLOAT)
    low = Column(FLOAT)
    close = Column(FLOAT)
    pre_close = Column(FLOAT)
    change = Column(FLOAT)
    pct_chg = Column(FLOAT)
    vol = Column(FLOAT)
    amount = Column(FLOAT)

    __table_args__ = (
        Index('idx_stockdaily_qfq', 'code', 'trade_date', unique=True),
    )


# 综合成长因子
class FactorComGrowth(Base):
    __tablename__ = 'factor_growth'

    id = Column(Integer, primary_key=True, autoincrement=True)
    code = Column(VARCHAR)
    end_date = Column(Date)
    f_ann_date = Column(Date)
    industry_code = Column(VARCHAR)
    signal_growth = Column(FLOAT)


"""
    宏观数据
"""


class MacroCPI(Base):
    __tablename__ = 'macro_cpi_cn'
    id = Column(Integer, primary_key=True, autoincrement=True)
    month = Column(VARCHAR)
    nt_val = Column(FLOAT)
    nt_yoy = Column(FLOAT)
    nt_mom = Column(FLOAT)
    nt_accu = Column(FLOAT)
    town_val = Column(FLOAT)
    town_yoy = Column(FLOAT)
    town_mom = Column(FLOAT)
    town_accu = Column(FLOAT)
    cnt_yoy = Column(FLOAT)
    cnt_mom = Column(FLOAT)
    cnt_accu = Column(FLOAT)


class MacroPPI(Base):
    __tablename__ = 'macro_ppi_cn'
    id = Column(Integer, primary_key=True, autoincrement=True)
    month = Column(VARCHAR)
    nt_val = Column(FLOAT)
    nt_yoy = Column(FLOAT)
    nt_mom = Column(FLOAT)
    nt_accu = Column(FLOAT)
    town_val = Column(FLOAT)
    town_yoy = Column(FLOAT)
    town_mom = Column(FLOAT)
    town_accu = Column(FLOAT)
    cnt_val = Column(FLOAT)
    cnt_yoy = Column(FLOAT)
    cnt_mom = Column(FLOAT)
    cnt_accu = Column(FLOAT)


class MacroSocialFinance(Base):
    __tablename__ = 'macro_social_financing_cn'
    id = Column(Integer, primary_key=True, autoincrement=True)
    month = Column(VARCHAR)
    inc_month = Column(FLOAT)
    inc_cumval = Column(FLOAT)
    stk_endval = Column(FLOAT)


class MacroBondCN(Base):
    __tablename__ = 'bond_term_cn'
    id = Column(Integer, primary_key=True, autoincrement=True)
    trade_date = Column(Date)
    ts_code = Column(VARCHAR)
    curve_name = Column(VARCHAR)
    curve_type = Column(VARCHAR)
    curve_term = Column(FLOAT)
    curve_yield = Column(FLOAT)
    yield_rate = Column(FLOAT)

    createtime = Column(TIMESTAMP, default=func.now())
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())


class MacroBondUS(Base):
    __tablename__ = 'bond_term_us'
    id = Column(Integer, primary_key=True, autoincrement=True)
    date = Column(Date)
    m1 = Column(FLOAT)
    m2 = Column(FLOAT)
    m3 = Column(FLOAT)
    m6 = Column(FLOAT)
    y1 = Column(FLOAT)
    y2 = Column(FLOAT)
    y3 = Column(FLOAT)
    y5 = Column(FLOAT)
    y7 = Column(FLOAT)
    y10 = Column(FLOAT)
    y20 = Column(FLOAT)
    y30 = Column(FLOAT)

    createtime = Column(TIMESTAMP, default=func.now())
    updatetime = Column(TIMESTAMP, default=func.now(), onupdate=func.now())


"""
    数据库初始化
"""

Base.metadata.create_all(engine)  # 本地数据库无表则创建

# 显式指定远程数据库检查表名称
remote_tables = [
    Table(FactorPeriodVolatility.__tablename__, Base.metadata),
    Table(FactorReturnWgt.__tablename__, Base.metadata),
    Table(FactorEarningSuprise.__tablename__, Base.metadata)
]

Base.metadata.create_all(bind=engine2, tables=remote_tables)  # 远程数据库无表则创建
